#include <rclcpp/rclcpp.hpp>
#include <geometry_msgs/msg/twist_stamped.hpp>
#include <geometry_msgs/msg/wrench_stamped.hpp>
#include <geometry_msgs/msg/pose_stamped.hpp>

#include <tf2_ros/transform_listener.h>
#include <tf2_ros/buffer.h>
#include <tf2/time.h>

#include <eigen3/Eigen/Dense>
#include <algorithm>
#include <string>
#include <vector>
#include <memory>
#include <cmath>

#include <moveit_servo/servo_parameters.h>
#include <moveit_servo/servo.h>
#include <moveit/planning_scene_monitor/planning_scene_monitor.h>
#include <moveit/move_group_interface/move_group_interface.h>

#include <std_srvs/srv/trigger.hpp>

using namespace std::chrono_literals;
static const rclcpp::Logger LOGGER = rclcpp::get_logger("force_servo_control_node");

class ForceServoControlNode : public rclcpp::Node
{
public:
  ForceServoControlNode() : Node("force_servo_control_node")
  {
    // —— 基本运动参数 ——
    vx_ = declare_parameter<double>("vx", 0.02);
    stop_x_threshold_ = declare_parameter<double>("stop_x", 0.08);
    publish_period_ms_ = declare_parameter<int>("period_ms", 5);  // 默认10是 100 Hz
    cmd_topic_ = declare_parameter<std::string>("cmd_topic", "force_servo_control_node/delta_twist_cmds");

    // —— 力控区间 ——
    F_low_  = declare_parameter<double>("F_low", 0.6);
    F_high_ = declare_parameter<double>("F_high", 0.7);
    v_z_max_= declare_parameter<double>("v_z_max", 0.05);
    down_sign_ = declare_parameter<double>("down_sign", -1.0);

    // —— PID 参数 ——
    kp_ = declare_parameter<double>("kp", 0.03);
    ki_ = declare_parameter<double>("ki", 0.05);
    kd_ = declare_parameter<double>("kd", 0.01);
    v_fast_down_ = declare_parameter<double>("v_fast_down", 0.02);
    F_touch_ = declare_parameter<double>("F_touch", 0.15);
    alpha_ = declare_parameter<double>("force_ema_alpha", 0.2);

    F_ref_ = 0.5 * (F_low_ + F_high_);
    tol_   = 0.5 * (F_high_ - F_low_);

    // —— 订阅力数据 ——
    force_sub_ = create_subscription<geometry_msgs::msg::WrenchStamped>(
        "/force_data", 50,
        std::bind(&ForceServoControlNode::force_callback, this, std::placeholders::_1));

    // —— TF & Publisher ——
    tf_buffer_ = std::make_shared<tf2_ros::Buffer>(this->get_clock());
    tf_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);
    twist_cmd_pub_ = create_publisher<geometry_msgs::msg::TwistStamped>(cmd_topic_, 10);

    // —— Reset 服务 ——
    reset_srv_ = this->create_service<std_srvs::srv::Trigger>(
        "reset",
        [this](const std::shared_ptr<std_srvs::srv::Trigger::Request> /*req*/,
               std::shared_ptr<std_srvs::srv::Trigger::Response> res) {
          goToTargetPose();          // 回到初始位姿
          state_ = State::RUN;       // 状态复位
          in_contact_ = false;
          integral_ = 0.0;
          prev_e_ = 0.0;
          stopped_published_ = false;
          res->success = true;
          res->message = "Force servo reset: moved to target pose and switched to RUN";
          RCLCPP_INFO(LOGGER, "Reset: moved to target pose, state=RUN");
        });

    // —— 参数动态回调 ——
    param_callback_handle_ = this->add_on_set_parameters_callback(
        std::bind(&ForceServoControlNode::param_callback, this, std::placeholders::_1));
  }

  void initialize()
  {
    setupPlanningSceneMonitor();
    setupServo();

    // 初始位姿移动（可选）
    goToTargetPose();

    state_ = State::RUN;
    stopped_published_ = false;

    // 定时控制循环
    timer_ = create_wall_timer(
        std::chrono::milliseconds(publish_period_ms_),
        std::bind(&ForceServoControlNode::control_loop, this));
  }

private:
  // —— 通信对象 ——  
  rclcpp::Subscription<geometry_msgs::msg::WrenchStamped>::SharedPtr force_sub_;
  rclcpp::Publisher<geometry_msgs::msg::TwistStamped>::SharedPtr twist_cmd_pub_;
  std::shared_ptr<tf2_ros::Buffer> tf_buffer_;
  std::shared_ptr<tf2_ros::TransformListener> tf_listener_;
  std::shared_ptr<planning_scene_monitor::PlanningSceneMonitor> planning_scene_monitor_;
  std::unique_ptr<moveit_servo::Servo> servo_;
  rclcpp::TimerBase::SharedPtr timer_;
  rclcpp::Service<std_srvs::srv::Trigger>::SharedPtr reset_srv_;
  OnSetParametersCallbackHandle::SharedPtr param_callback_handle_;

  // —— 状态 & 变量 ——  
  enum class State { RUN, STOPPED };
  State state_{State::RUN};

  bool stopped_published_{false};  // STOP时是否已经发过一次0

  Eigen::Vector3d current_force_{Eigen::Vector3d::Zero()};
  double Fx_f_{0.0}, Fz_f_{0.0};

  // —— 参数 ——  
  double vx_, stop_x_threshold_;
  int publish_period_ms_;
  std::string cmd_topic_;
  double F_low_, F_high_, v_z_max_, down_sign_;
  std::vector<double> target_position_{-0.021511, 0.0, 0.09558};
  std::vector<double> target_orientation_{0.0, 0.42257, 0.0, 0.90633};

  // —— PID & 控制 ——  
  double kp_, ki_, kd_;
  double v_fast_down_, F_touch_, alpha_;
  double F_ref_, tol_;
  bool in_contact_{false};
  double integral_{0.0};
  double prev_e_{0.0};

  // 参数动态回调
  rcl_interfaces::msg::SetParametersResult
  param_callback(const std::vector<rclcpp::Parameter> &params)
  {
    for (const auto &param : params) {
      if (param.get_name() == "kp") kp_ = param.as_double();
      else if (param.get_name() == "ki") ki_ = param.as_double();
      else if (param.get_name() == "kd") kd_ = param.as_double();
      else if (param.get_name() == "force_ema_alpha") alpha_ = param.as_double();
      else if (param.get_name() == "v_fast_down") v_fast_down_ = param.as_double();
    }
    RCLCPP_INFO(LOGGER, "Updated PID: kp=%.3f ki=%.3f kd=%.3f alpha=%.2f v_fast_down=%.3f",
                kp_, ki_, kd_, alpha_, v_fast_down_);
    rcl_interfaces::msg::SetParametersResult result;
    result.successful = true;
    return result;
  }

  void force_callback(const geometry_msgs::msg::WrenchStamped::SharedPtr msg)
  {
    Fx_f_ = alpha_ * Fx_f_ + (1.0 - alpha_) * msg->wrench.force.x;
    Fz_f_ = alpha_ * Fz_f_ + (1.0 - alpha_) * msg->wrench.force.z;
    current_force_.x() = Fx_f_;
    current_force_.z() = Fz_f_;
  }

  void goToTargetPose()
  {
    moveit::planning_interface::MoveGroupInterface move_group(shared_from_this(), "zc4");
    geometry_msgs::msg::Pose target_pose;
    target_pose.position.x = target_position_.at(0);
    target_pose.position.y = target_position_.at(1);
    target_pose.position.z = target_position_.at(2);
    target_pose.orientation.x = target_orientation_.at(0);
    target_pose.orientation.y = target_orientation_.at(1);
    target_pose.orientation.z = target_orientation_.at(2);
    target_pose.orientation.w = target_orientation_.at(3);
    move_group.setPoseTarget(target_pose);
    moveit::planning_interface::MoveGroupInterface::Plan my_plan;
    if (move_group.plan(my_plan) == moveit::planning_interface::MoveItErrorCode::SUCCESS) {
      move_group.move();
    }
  }

  void control_loop()
  {
    // 获取末端位置
    double cur_x = 0.0;
    try {
      auto tf = tf_buffer_->lookupTransform("base_link", "end", tf2::TimePointZero);
      cur_x = tf.transform.translation.x;
    } catch (...) { return; }

    auto msg = std::make_unique<geometry_msgs::msg::TwistStamped>();
    msg->header.stamp = this->now();
    msg->header.frame_id = "base_link";

    // Stop 判据
    if (state_ == State::RUN && cur_x >= stop_x_threshold_) {
      state_ = State::STOPPED;
    }

    if (state_ == State::STOPPED) {
      if (!stopped_published_) {
        msg->twist.linear.x = 0.0;
        msg->twist.linear.y = 0.0;
        msg->twist.linear.z = 0.0;
        msg->twist.angular.x = 0.0;
        msg->twist.angular.y = 0.0;
        msg->twist.angular.z = 0.0;
        twist_cmd_pub_->publish(*msg);
        RCLCPP_INFO(LOGGER, "STOPPED: published zero twist once");
        stopped_published_ = true;
      }
      return;
    }

    // RUN 状态：PID 力控
    const double Fx = current_force_.x();
    const double Fz = current_force_.z();
    const double F_mag = std::hypot(Fx, Fz);
    const double dt = static_cast<double>(publish_period_ms_) / 1000.0;

    if (!in_contact_ && F_mag >= F_touch_) {
      in_contact_ = true;
      prev_e_ = (F_ref_ - F_mag);
    }

    double e = F_ref_ - F_mag;
    double v_z_cmd = 0.0;

    if (!in_contact_) {
      v_z_cmd = down_sign_ * v_fast_down_;
    } else {
      if (std::fabs(e) <= tol_) {
        integral_ *= 0.95;
        v_z_cmd = 0.0;
      } else {
        integral_ += e * dt;
        const double I_MAX = 0.8 * (v_z_max_ / std::max(1e-6, ki_));
        integral_ = std::clamp(integral_, -I_MAX, I_MAX);
        double de = (e - prev_e_) / std::max(1e-6, dt);
        prev_e_ = e;
        double u = kp_ * e + ki_ * integral_ + kd_ * de;
        v_z_cmd = down_sign_ * u;
      }
    }

    v_z_cmd = std::clamp(v_z_cmd, -v_z_max_, v_z_max_);

    msg->twist.linear.x = vx_;
    msg->twist.linear.y = 0.0;
    msg->twist.linear.z = v_z_cmd;
    twist_cmd_pub_->publish(*msg);
  }

  void setupPlanningSceneMonitor()
  {
    planning_scene_monitor_ = std::make_shared<planning_scene_monitor::PlanningSceneMonitor>(
        shared_from_this(), "robot_description", tf_buffer_, "planning_scene_monitor");
    if (planning_scene_monitor_->getPlanningScene()) {
      planning_scene_monitor_->startStateMonitor("/joint_states");
      planning_scene_monitor_->setPlanningScenePublishingFrequency(25);
      planning_scene_monitor_->startPublishingPlanningScene(
          planning_scene_monitor::PlanningSceneMonitor::UPDATE_SCENE,
          "/moveit_servo/publish_planning_scene");
      planning_scene_monitor_->startSceneMonitor();
      planning_scene_monitor_->providePlanningSceneService();
    }
  }

  void setupServo()
  {
    auto servo_parameters = moveit_servo::ServoParameters::makeServoParameters(shared_from_this());
    if (!servo_parameters) throw std::runtime_error("Failed to load servo parameters");
    servo_ = std::make_unique<moveit_servo::Servo>(shared_from_this(), servo_parameters, planning_scene_monitor_);
    servo_->start();
  }
};

int main(int argc, char **argv)
{
  rclcpp::init(argc, argv);
  auto node = std::make_shared<ForceServoControlNode>();
  node->initialize();
  rclcpp::spin(node);
  rclcpp::shutdown();
  return 0;
}
